_base_ = ["stage2.py"]

# Define model components
model = dict(cond_embed=True)
grad_ckpt_buffer_size = 25 * 1024**3

condition_config = dict(
    t2v=1,
    i2v_head=5,
    i2v_loop=1,
    i2v_tail=1,
)
is_causal_vae = True

bucket_config = {
    "_delete_": True,
    "256px": {
        1: (1.0, 195),
        5: (1.0, 80),
        9: (1.0, 80),
        13: (1.0, 80),
        17: (1.0, 80),
        21: (1.0, 80),
        25: (1.0, 80),
        29: (1.0, 80),
        33: (1.0, 80),
        37: (1.0, 40),
        41: (1.0, 40),
        45: (1.0, 40),
        49: (1.0, 40),
        53: (1.0, 40),
        57: (1.0, 40),
        61: (1.0, 40),
        65: (1.0, 40),
        69: (1.0, 28),
        73: (1.0, 28),
        77: (1.0, 28),
        81: (1.0, 28),
        85: (1.0, 28),
        89: (1.0, 28),
        93: (1.0, 28),
        97: (1.0, 28),
        101: (1.0, 23),
        105: (1.0, 23),
        109: (1.0, 23),
        113: (1.0, 23),
        117: (1.0, 23),
        121: (1.0, 23),
        125: (1.0, 23),
        129: (1.0, 23),
    },
    "768px": {
        1: (0.5, 38),
        5: (0.5, 10),
        9: (0.5, 10),
        13: (0.5, 10),
        17: (0.5, 10),
        21: (0.5, 10),
        25: (0.5, 10),
        29: (0.5, 10),
        33: (0.5, 10),
        37: (0.5, 5),
        41: (0.5, 5),
        45: (0.5, 5),
        49: (0.5, 5),
        53: (0.5, 5),
        57: (0.5, 5),
        61: (0.5, 5),
        65: (0.5, 5),
        69: (0.5, 3),
        73: (0.5, 3),
        77: (0.5, 3),
        81: (0.5, 3),
        85: (0.5, 3),
        89: (0.5, 3),
        93: (0.5, 3),
        97: (0.5, 3),
        101: (0.5, 2),
        105: (0.5, 2),
        109: (0.5, 2),
        113: (0.5, 2),
        117: (0.5, 2),
        121: (0.5, 2),
        125: (0.5, 2),
        129: (0.5, 2),
    },
}
